// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use paste::paste;

use crate::bincode::Encoder;

pub trait Encode {
    unsafe fn encode_unchecked<E: Encoder>(&self, encoder: &mut E, buf: &mut [u8]);
    fn encode<E: Encoder>(&self, encoder: &mut E, buf: &mut Vec<u8>) -> Result<(), E::Error>;
}

macro_rules! impl_encode_for_builtin {
    ($type:ident) => {
        impl Encode for $type {
            unsafe fn encode_unchecked<E: Encoder>(&self, encoder: &mut E, buf: &mut [u8]) {
                unsafe { paste! { encoder.[<put_ $type _unchecked>](buf, *self) } }
            }
            
            fn encode<E: Encoder>(&self, encoder: &mut E, buf: &mut Vec<u8>) -> Result<(), E::Error> {
                paste! { encoder.[<put_ $type>](buf, *self) }
            }
        }
    };
}

impl_encode_for_builtin!(bool);
impl_encode_for_builtin!(u8);
impl_encode_for_builtin!(u16);
impl_encode_for_builtin!(u32);
impl_encode_for_builtin!(u64);
impl_encode_for_builtin!(i8);
impl_encode_for_builtin!(i16);
impl_encode_for_builtin!(i32);
impl_encode_for_builtin!(i64);
impl_encode_for_builtin!(f32);
impl_encode_for_builtin!(f64);

impl Encode for &str {
    /// # Safety
    /// 
    /// This method is unsafe because it allows the caller to provide a buffer with insufficient capacity.
    /// 
    /// # Examples
    /// 
    /// ```
    /// use ohmw_core::bincode::{Encode, DefaultEncoder};
    /// 
    /// let mut buf = vec![0; 17];
    /// let val: &str = "Hello, world!";
    /// unsafe { val.encode_unchecked(&mut DefaultEncoder::new(), &mut buf) };
    /// assert_eq!(buf, vec![13, 0, 0, 0, 72, 101, 108, 108, 111, 44, 32, 119, 111, 114, 108, 100, 33]);
    /// ```
    unsafe fn encode_unchecked<E: Encoder>(&self, encoder: &mut E, buf: &mut [u8]) {
        unsafe { self.as_bytes().encode_unchecked(encoder, buf) }
    }

    /// # Examples
    /// 
    /// ```
    /// use ohmw_core::bincode::{Encode, DefaultEncoder};
    /// 
    /// let mut buf = Vec::new();
    /// let val: &str = "Hello, world!";
    /// val.encode(&mut DefaultEncoder::new(), &mut buf)
    ///     .expect("failed to encode message");
    /// assert_eq!(buf, vec![13, 0, 0, 0, 72, 101, 108, 108, 111, 44, 32, 119, 111, 114, 108, 100, 33]);
    /// ```
    fn encode<E: Encoder>(&self, encoder: &mut E, buf: &mut Vec<u8>) -> Result<(), E::Error> {
        self.as_bytes().encode(encoder, buf)
    }
}

impl Encode for String {
    /// # Safety
    /// 
    /// This method is unsafe because it allows the caller to provide a buffer with insufficient capacity.
    /// 
    /// # Examples
    /// 
    /// ```
    /// use ohmw_core::bincode::{Encode, DefaultEncoder};
    /// 
    /// let mut buf = vec![0; 17];
    /// let val: String = "Hello, world!".to_string();
    /// unsafe { val.encode_unchecked(&mut DefaultEncoder::new(), &mut buf) };
    /// assert_eq!(buf, vec![13, 0, 0, 0, 72, 101, 108, 108, 111, 44, 32, 119, 111, 114, 108, 100, 33]);
    /// ```
    unsafe fn encode_unchecked<E: Encoder>(&self, encoder: &mut E, buf: &mut [u8]) {
        unsafe { self.as_str().encode_unchecked(encoder, buf) }
    }

    /// # Examples
    /// 
    /// ```
    /// use ohmw_core::bincode::{Encode, DefaultEncoder};
    /// 
    /// let mut buf = Vec::new();
    /// let val: String = "Hello, world!".to_string();
    /// val.encode(&mut DefaultEncoder::new(), &mut buf)
    ///     .expect("failed to encode message");
    /// assert_eq!(buf, vec![13, 0, 0, 0, 72, 101, 108, 108, 111, 44, 32, 119, 111, 114, 108, 100, 33]);
    /// ```
    fn encode<E: Encoder>(&self, encoder: &mut E, buf: &mut Vec<u8>) -> Result<(), E::Error> {
        self.as_str().encode(encoder, buf)
    }
}

impl<T: Encode> Encode for &[T] {
    unsafe fn encode_unchecked<E: Encoder>(&self, encoder: &mut E, buf: &mut [u8]) {
        unsafe { encoder.put_u32_unchecked(buf, self.len() as u32) };
        if unty::type_equal::<T, u8>() {
            unsafe { encoder.put_bytes_unchecked(buf, std::mem::transmute(*self)) };
        } else {
            for item in self.iter() {
                unsafe { item.encode_unchecked(encoder, buf) };
            }
        }
    }

    fn encode<E: Encoder>(&self, encoder: &mut E, buf: &mut Vec<u8>) -> Result<(), E::Error> {
        if unty::type_equal::<T, u8>() {
            encoder.ensure_capacity(buf, size_of::<u32>() + self.len())?;
            unsafe { encoder.put_u32_unchecked(buf, self.len() as u32) };
            unsafe { encoder.put_bytes_unchecked(buf, std::mem::transmute(*self)) };
            Ok(())
        } else {
            (self.len() as u32).encode(encoder, buf)?;
            for item in self.iter() {
                item.encode(encoder, buf)?;
            }
            Ok(())
        }
    }
}

impl<T: Encode, const N: usize> Encode for [T; N] {
    /// # Safety
    /// 
    /// This method is unsafe because it allows the caller to provide a buffer with insufficient capacity.
    /// 
    /// # Examples
    /// 
    /// ```
    /// use ohmw_core::bincode::{Encode, DefaultEncoder};
    /// 
    /// let mut buf = vec![0; 5];
    /// let arr: [u8; 5] = [10, 20, 30, 40, 50];
    /// unsafe { arr.encode_unchecked(&mut DefaultEncoder::new(), &mut buf) };
    /// assert_eq!(buf, vec![10, 20, 30, 40, 50]);
    /// ```
    unsafe fn encode_unchecked<E: Encoder>(&self, encoder: &mut E, buf: &mut [u8]) {
        if unty::type_equal::<T, u8>() {
            unsafe { encoder.put_bytes_unchecked(buf, std::mem::transmute(self.as_slice())) };
        } else {
            for item in self.iter() {
                unsafe { item.encode_unchecked(encoder, buf) };
            }
        }
    }

    /// # Examples
    /// 
    /// ```
    /// use ohmw_core::bincode::{Encode, DefaultEncoder};
    /// 
    /// let mut buf = Vec::new();
    /// let arr: [u8; 5] = [10, 20, 30, 40, 50];
    /// arr.encode(&mut DefaultEncoder::new(), &mut buf)
    ///     .expect("failed to encode message");
    /// assert_eq!(buf, vec![10, 20, 30, 40, 50]);
    /// ```
    fn encode<E: Encoder>(&self, encoder: &mut E, buf: &mut Vec<u8>) -> Result<(), E::Error> {
        if unty::type_equal::<T, u8>() {
            encoder.ensure_capacity(buf, N)?;
            unsafe { encoder.put_bytes_unchecked(buf, std::mem::transmute(self.as_slice())) };
            Ok(())
        } else {
            for item in self.iter() {
                item.encode(encoder, buf)?;
            }
            Ok(())
        }
    }
}

impl<T: Encode> Encode for Vec<T> {
    unsafe fn encode_unchecked<E: Encoder>(&self, encoder: &mut E, buf: &mut [u8]) {
        unsafe { self.as_slice().encode_unchecked(encoder, buf) }
    }

    fn encode<E: Encoder>(&self, encoder: &mut E, buf: &mut Vec<u8>) -> Result<(), E::Error> {
        self.as_slice().encode(encoder, buf)
    }
}

impl<T: Encode> Encode for Option<T> {
    unsafe fn encode_unchecked<E: Encoder>(&self, encoder: &mut E, buf: &mut [u8]) {
        match self {
            Some(val) => {
                unsafe { encoder.put_bool_unchecked(buf, true) };
                unsafe { val.encode_unchecked(encoder, buf) }
            },
            None => unsafe { encoder.put_bool_unchecked(buf, false) },
        }
    }

    fn encode<E: Encoder>(&self, encoder: &mut E, buf: &mut Vec<u8>) -> Result<(), E::Error> {
        match self {
            Some(val) => {
                encoder.put_bool(buf, true)?;
                val.encode(encoder, buf)
            },
            None => {
                encoder.put_bool(buf, false)?;
                Ok(())
            },
        }
    }
}
